/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2022-2022.  All rights reserved.2023.
 */

package service

import (
	"bytes"
	"crypto"
	"crypto/rsa"
	"crypto/tls"
	"crypto/x509"
	"encoding/base64"
	"encoding/pem"
	"encoding/xml"
	"fmt"
	"io/ioutil"
	"net/http"
	"net/url"
	"os"
	"prim-server/common"
	"prim-server/config"
	"prim-server/logger"
	"regexp"
	"strings"
	"time"

	"github.com/beevik/etree"
	"github.com/crewjam/saml"
	"github.com/crewjam/saml/xmlenc"
	"github.com/google/uuid"
	"github.com/pkg/errors"
	dsig "github.com/russellhaering/goxmldsig"
	"github.com/ucarion/c14n"
)

const (
	SAML_REQUEST = "SAMLRequest"
	RELAY_STATE  = "RelayState"
	SIGALG       = "SigAlg"
	SIGNATURE    = "Signature"

	DEFAULT_ASSERTION_VALID_TIME_MIN = 5
)

var spMeta *saml.EntityDescriptor
var idpMeta *saml.EntityDescriptor

var (
	bpId = config.GetConfig().BSSConfig.BPID
)

func readMetaData(metadataPath string, metadata *saml.EntityDescriptor) error {
	//读取sp metadata的内容
	file, err := os.Open(metadataPath)
	if err != nil {
		logger.Error("read MetaData error: %v, path: %s", err, metadata)
		return err
	}
	defer file.Close()
	buffer, err := ioutil.ReadAll(file)
	if err != nil {
		logger.Error("read MetaData error: %v, path: %s", err, metadata)
		return err
	}
	err = xml.Unmarshal(buffer, metadata)
	if err != nil {
		logger.Error("SAMLLogin error: %v, path: %s", err, metadata)
		return err
	}
	return nil
}
func GetSpMetaData() (*saml.EntityDescriptor, error) {
	if spMeta == nil {
		spMeta = new(saml.EntityDescriptor)
		err := readMetaData(config.GetConfig().SMALConfig.SPMetaDataPath, spMeta)
		if err != nil {
			return nil, err
		}
	}
	return spMeta, nil
}
func GetIdpMetaData() (*saml.EntityDescriptor, error) {
	if idpMeta == nil {
		idpMeta = new(saml.EntityDescriptor)
		err := readMetaData(config.GetConfig().SMALConfig.IDPMetaDataPath, idpMeta)
		if err != nil {
			return nil, err
		}
	}
	return idpMeta, nil
}
func VerifySignature(spMetadata *saml.EntityDescriptor, query common.SAMLRequest) error {
	//公钥证书
	publicKey := spMetadata.SPSSODescriptors[0].KeyDescriptors[0].KeyInfo.X509Data.X509Certificates[0].Data
	publicKey = regexp.MustCompile(`\s+`).ReplaceAllString(publicKey, "")
	publicKeyBytes, err := base64.StdEncoding.DecodeString(publicKey)
	if err != nil {
		logger.Error("verify signature, decode string error: %v, sinature: %s", err, query.Signature)
		return err
	}
	parseKey, err := x509.ParseCertificate(publicKeyBytes)
	if err != nil {
		logger.Error("verify signature, parse certificate error: %v, sinature: %s", err, query.Signature)
		return err
	}

	//解析Signature
	decodeSignature, err := base64.StdEncoding.DecodeString(query.Signature)
	if err != nil {
		logger.Error("verify signature, decode string error: %v, sinature: %s", err, query.Signature)
		return err
	}
	//按顺序拼接SAMLRequest, RelayState, SigAlg
	_, signedData := concat(query.SAMLRequest, query.RelayState, query.SigAlg)
	signatureAlgorithm, err := extractSignatureAlgorithm(query.SigAlg)
	if err != nil {
		logger.Error("SAMLLogin error: %v", err)
		return err
	}
	//计算hash值
	hashAlgorithm := getHashAlgorithm(signatureAlgorithm)
	hash := crypto.Hash.New(hashAlgorithm)
	if _, err = hash.Write(signedData); err != nil {
		logger.Error("SAMLLogin error: %v", err)
		return err
	}
	hashedData := hash.Sum(nil)

	//验证签名
	rsaKey, ok := parseKey.PublicKey.(*rsa.PublicKey)
	if !ok {
		return errors.New("cannot convet")
	}
	err = rsa.VerifyPKCS1v15(rsaKey, hashAlgorithm, hashedData, decodeSignature)
	if err != nil {
		logger.Warn("SAMLLogin error: %v", err)
		return err
	}
	return nil
}
func BuildSamlResponse(samlRequest saml.AuthnRequest, spMetadata *saml.EntityDescriptor, userId string) (*etree.Element, error) {
	samlAssertion, err := buildAssertion(userId, samlRequest.AssertionConsumerServiceURL, samlRequest.ID, spMetadata)
	if err != nil {
		return nil, err
	}
	//签名assertion
	assertionEl, err := SignedAssertion(&samlAssertion)
	if err != nil {
		logger.Error("signed assertion error: %v", err)
		return nil, err
	}
	idpMeta, err := GetIdpMetaData()
	if err != nil {
		return nil, err
	}
	consent := "urn:oasis:names:tc:SAML:2.0:consent:unspecified"
	samlResponse := &saml.Response{
		ID:           "_" + uuid.New().String(),
		Issuer:       buildIssuer(idpMeta.EntityID),
		Status:       buildStatus(saml.StatusSuccess),
		IssueInstant: time.Now().UTC(),
		InResponseTo: samlRequest.ID,
		Consent:      consent,
		Destination:  samlRequest.AssertionConsumerServiceURL,
		Version:      "2.0",
	}
	responseEl, err := SignedResponse(samlResponse, assertionEl)
	if err != nil {
		return nil, err
	}
	return responseEl, nil
}
func SignedResponse(response *saml.Response, assertionEl *etree.Element) (*etree.Element, error) {
	responseEl := response.Element()
	responseEl.AddChild(assertionEl)
	// Sign the response element (we've already signed the Assertion element)
	// 读取私钥
	privateFileBuffer, err := ioutil.ReadFile(config.GetConfig().SMALConfig.SignaturePrivateKeyPath)
	if err != nil {
		logger.Error("read file error: %v", err)
		return nil, err
	}
	block, _ := pem.Decode(privateFileBuffer)
	privatekey, err := x509.ParsePKCS1PrivateKey(block.Bytes)
	if err != nil {
		logger.Error("parse certificate error: %v", err)
		return nil, err
	}
	//读取公钥
	publicFileBuffer, err := ioutil.ReadFile(config.GetConfig().SMALConfig.SignatureCAKeyPath)
	if err != nil {
		logger.Error("read file error: %v", err)
		return nil, err
	}
	publicblock, _ := pem.Decode(publicFileBuffer)
	publickey, err := x509.ParseCertificate(publicblock.Bytes)
	if err != nil {
		logger.Error("parse certificate error: %v", err)
		return nil, err
	}
	keyPair := tls.Certificate{
		Certificate: [][]byte{publickey.Raw},
		PrivateKey:  privatekey,
		Leaf:        publickey,
	}

	keyStore := dsig.TLSCertKeyStore(keyPair)

	signatureMethod := dsig.RSASHA256SignatureMethod

	canonicalizerPrefixList := ""

	signingContext := dsig.NewDefaultSigningContext(keyStore)
	signingContext.Canonicalizer = dsig.MakeC14N10ExclusiveCanonicalizerWithPrefixList(canonicalizerPrefixList)
	if err := signingContext.SetSignatureMethod(signatureMethod); err != nil {
		return nil, err
	}

	signedResponseEl, err := signingContext.SignEnveloped(responseEl)
	if err != nil {
		return nil, err
	}

	sigEl := signedResponseEl.ChildElements()[len(signedResponseEl.ChildElements())-1]
	response.Signature = sigEl
	responseEl = response.Element()
	responseEl.AddChild(assertionEl)
	return responseEl, nil
}
func SignedAssertion(samlAssertion *saml.Assertion) (*etree.Element, error) {
	assertionEl, err := MakeAssertionEl(samlAssertion)
	if err != nil {
		return nil, err
	}
	return assertionEl, nil
}

// 读取私钥文件
func readSignaturePrivateKey() (*rsa.PrivateKey, error) {
	// 读取私钥
	privateFileBuffer, err := ioutil.ReadFile(config.GetConfig().SMALConfig.SignaturePrivateKeyPath)
	if err != nil {
		logger.Error("read signature privateKey error: %v", err)
		return nil, err
	}
	block, _ := pem.Decode(privateFileBuffer)
	privatekey, err := x509.ParsePKCS1PrivateKey(block.Bytes)
	if err != nil {
		logger.Error("read signature privateKey error: %v", err)
		return nil, err
	}
	return privatekey, nil
}

// 读取公钥
func readSignaturePublicKey() (*x509.Certificate, error) {
	// 读取私钥
	publicFileBuffer, err := ioutil.ReadFile(config.GetConfig().SMALConfig.SignatureCAKeyPath)
	if err != nil {
		logger.Error("read signature publicKey error: %v", err)
		return nil, err
	}
	publicblock, _ := pem.Decode(publicFileBuffer)
	publickey, err := x509.ParseCertificate(publicblock.Bytes)
	if err != nil {
		logger.Error("read signature publicKey error: %v", err)
		return nil, err
	}
	return publickey, nil
}

// 读取sp metadata中的公钥
func readCertInSpMetadata() (*x509.Certificate, error) {
	spMeta, err := GetSpMetaData()
	if err != nil {
		return nil, err
	}
	certStr := spMeta.SPSSODescriptors[0].KeyDescriptors[0].KeyInfo.X509Data.X509Certificates[0].Data
	certStr = regexp.MustCompile(`\s+`).ReplaceAllString(certStr, "")
	certBytes, err := base64.StdEncoding.DecodeString(certStr)
	if err != nil {
		return nil, errors.Errorf("cannot decode certificate base64: %v", err)
	}
	certBuf, err := x509.ParseCertificate(certBytes)
	if err != nil {
		return nil, errors.Errorf("cannot parse certificate: %v", err)
	}
	return certBuf, nil
}
func makeC14(keyStore dsig.TLSCertKeyStore, assertion *saml.Assertion) (retAssertionEl, retSignedAssertionEl *etree.Element, err error) {
	signatureMethod := dsig.RSASHA256SignatureMethod

	signingContext := dsig.NewDefaultSigningContext(keyStore)
	signingContext.Canonicalizer = dsig.MakeC14N10ExclusiveCanonicalizerWithPrefixList("")
	if err := signingContext.SetSignatureMethod(signatureMethod); err != nil {
		return nil, nil, err
	}
	assertionEl := assertion.Element()

	signedAssertionEl, err := signingContext.SignEnveloped(assertionEl)
	if err != nil {
		return nil, nil, err
	}

	sigEl := signedAssertionEl.Child[len(signedAssertionEl.Child)-1]
	convertSignature, ok := sigEl.(*etree.Element)
	if !ok {
		return nil, nil, errors.New("sigEl convert fail")
	}
	assertion.Signature = convertSignature

	return assertionEl, signedAssertionEl, nil
}
func MakeAssertionEl(assertion *saml.Assertion) (*etree.Element, error) {
	//读取私钥
	privatekey, err := readSignaturePrivateKey()
	if err != nil {
		logger.Error("make assertionEl error: %v", err)
		return nil, err
	}
	//读取公钥
	publickey, err := readSignaturePublicKey()
	if err != nil {
		logger.Error("make assertionEl error: %v", err)
		return nil, err
	}
	keyPair := tls.Certificate{
		Certificate: [][]byte{publickey.Raw},
		PrivateKey:  privatekey,
		Leaf:        publickey,
	}
	keyStore := dsig.TLSCertKeyStore(keyPair)

	assertionEl, signedAssertionEl, err := makeC14(keyStore, assertion)
	if err != nil {
		return nil, err
	}

	signedAssertionEl = assertion.Element()

	assertionEl = signedAssertionEl

	var signedAssertionBuf []byte
	{
		doc := etree.NewDocument()
		doc.SetRoot(signedAssertionEl)
		signedAssertionBuf, err = doc.WriteToBytes()
		if err != nil {
			return nil, err
		}
	}
	certBuf, err := readCertInSpMetadata()
	if err != nil {
		return nil, err
	}
	encryptor := xmlenc.OAEP()
	encryptor.BlockCipher = xmlenc.AES128CBC
	encryptor.DigestMethod = &xmlenc.SHA1
	encryptedDataEl, err := encryptor.Encrypt(certBuf, signedAssertionBuf, nil)
	if err != nil {
		return nil, err
	}
	encryptedDataEl.CreateAttr("Type", "http://www.w3.org/2001/04/xmlenc#Element")

	encryptedAssertionEl := etree.NewElement("saml:EncryptedAssertion")
	encryptedAssertionEl.AddChild(encryptedDataEl)
	assertionEl = encryptedAssertionEl

	return assertionEl, nil
}
func concat(samlRequest, relayState, sigAlg string) (string, []byte) {
	queryString := fmt.Sprintf("SAMLRequest=%s&RelayState=%s&SigAlg=%s",
		samlRequest, relayState, sigAlg)
	queryBytes := []byte(queryString)
	return queryString, queryBytes
}
func extractSignatureAlgorithm(sigAlg string) (string, error) {
	ss := strings.Split(sigAlg, "#")
	if len(ss) < 2 {
		return "", errors.New("unsupporsted SigAlgUrl: " + sigAlg)
	}
	if ss[1] == "rsa-sha1" {
		return "rsa-sha1", nil
	} else if ss[1] == "rsa-sha256" {
		return "rsa-sha256", nil
	} else {
		return "", errors.New("unsupporsted SigAlgUrl: " + sigAlg)
	}
}

func getHashAlgorithm(signatureAlgorithm string) crypto.Hash {
	switch signatureAlgorithm {
	case "rsa-sha1":
		return crypto.SHA1
	case "rsa-sha256":
		return crypto.SHA256
	default:
		return crypto.SHA1
	}
}
func buildSubject(subjectNameId, nameQualifier, recipient, inResponseTo string, notOnOrAfter time.Time) *saml.Subject {
	subject := &saml.Subject{
		NameID: &saml.NameID{
			Value:         subjectNameId,
			Format:        string(saml.TransientNameIDFormat),
			NameQualifier: nameQualifier,
		},
		SubjectConfirmations: []saml.SubjectConfirmation{
			{
				SubjectConfirmationData: &saml.SubjectConfirmationData{
					Recipient:    recipient,
					NotOnOrAfter: notOnOrAfter,
					InResponseTo: inResponseTo,
				},
			},
		},
	}
	return subject
}
func buildIssuer(issuerValue string) *saml.Issuer {
	issuer := &saml.Issuer{
		Value:  issuerValue,
		Format: "urn:oasis:names:tc:SAML:2.0:nameid-format:entity",
	}
	return issuer
}
func buildAssertion(userId, assertionConsumerUrl, requestId string, spMetadata *saml.EntityDescriptor) (saml.Assertion, error) {
	validateTime := time.Minute * time.Duration(DEFAULT_ASSERTION_VALID_TIME_MIN)
	idpMetaData, err := GetIdpMetaData()
	if err != nil {
		return saml.Assertion{}, err
	}
	entityId := idpMetaData.EntityID
	notBefore := time.Now().UTC()
	notOnOrAfter := notBefore.Add(validateTime)
	spEntityId := spMetadata.EntityID
	subject := buildSubject(entityId, spEntityId, assertionConsumerUrl, requestId, notOnOrAfter)
	issuer := buildIssuer(entityId)
	audienceRestriction := saml.AudienceRestriction{}
	audienceRestriction.Audience = saml.Audience{
		Value: spMetadata.EntityID,
	}
	conditions := &saml.Conditions{}
	conditions.AudienceRestrictions = append(conditions.AudienceRestrictions, audienceRestriction)
	conditions.NotBefore = notBefore
	conditions.NotOnOrAfter = notOnOrAfter
	authnStatement := saml.AuthnStatement{
		AuthnInstant: time.Now().UTC(),
		SessionIndex: uuid.New().String(),
		AuthnContext: saml.AuthnContext{},
	}
	pwdTransport := "urn:oasis:names:tc:SAML:2.0:ac:classes:PasswordProtectedTransport"
	authnStatement.AuthnContext.AuthnContextClassRef = &saml.AuthnContextClassRef{
		Value: pwdTransport,
	}

	assertion := saml.Assertion{
		ID:           uuid.New().String(),
		IssueInstant: time.Now().UTC(),
		Subject:      subject,
		Issuer:       *issuer,
		Conditions:   conditions,
		Version:      "2.0",
	}
	assertion.AuthnStatements = append(assertion.AuthnStatements, authnStatement)

	attribute := make(map[string][]string)
	attribute = buildAttribute(attribute, userId)
	attributeStatement := buildAttributeStatement(attribute)
	assertion.AttributeStatements = append(assertion.AttributeStatements, *attributeStatement)
	return assertion, nil
}
func buildAttribute(attribute map[string][]string, userId string) map[string][]string {
	//设置attributeValue
	attribute["xAccountId"] = []string{userId}
	attribute["xUserId"] = []string{userId}
	attribute["bpId"] = []string{bpId}
	return attribute
}
func buildAttributeStatement(attribute map[string][]string) *saml.AttributeStatement {
	attributeStatemnet := &saml.AttributeStatement{}
	attributeNameFormat := "urn:oasis:names:tc:SAML:2.0:attrname-format:uri"
	for name, values := range attribute {
		attr := saml.Attribute{
			FriendlyName: name,
			Name:         name,
			NameFormat:   attributeNameFormat,
		}
		for _, value := range values {
			attrValue := saml.AttributeValue{
				Type:  "xs:string",
				Value: value,
			}
			attr.Values = append(attr.Values, attrValue)
		}
		attributeStatemnet.Attributes = append(attributeStatemnet.Attributes, attr)
	}
	return attributeStatemnet
}
func buildStatus(value string) saml.Status {
	status := saml.Status{
		StatusCode: saml.StatusCode{
			Value: value,
		},
		StatusMessage: &saml.StatusMessage{
			Value: value,
		},
	}
	return status
}

func canonicalize(b []byte) []byte {
	decoder := xml.NewDecoder(bytes.NewReader(b))
	out, err := c14n.Canonicalize(decoder)
	if err != nil {
		logger.Error("canonicalize error: %v", err)
		return nil
	}
	return out
}
func sendSAMLResponse(assertionConsumerUrl, relayState string, samlResponse saml.Response) {
	//post请求应答路径
	samlResponseBytes, err := xml.Marshal(samlResponse)
	if err != nil {
		logger.Error("sendSAMLResponse error: %v", err)
		return
	}
	postResponse := base64.StdEncoding.EncodeToString(samlResponseBytes)

	formValues := url.Values{}
	formValues.Set("SAMLResponse", postResponse)
	formValues.Set("RelayState", relayState)
	formDataStr := formValues.Encode()
	formDataBytes := []byte(formDataStr)
	formBytesReader := bytes.NewReader(formDataBytes)
	client := &http.Client{}
	req, err := http.NewRequest("POST", assertionConsumerUrl, formBytesReader)
	if err != nil {
		logger.Error("sendSAMLResponse error: %v", err)
		return
	}
	req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	rest, err := client.Do(req)
	if err != nil {
		logger.Error("sendSAMLResponse error: %v", err)
		return
	}
	fmt.Printf("%v", rest)
	return
}
